# coding = utf-8

'''
保留只有肿瘤的数据
'''

import os
import numpy as np
import cv2
from scipy import ndimage
from skimage import measure
import prettytable as pt
from pathlib2 import Path
import math

def save_only_tumor():
    data_path = "/datasets/DongbeiDaxue/chengkun_remove"
    save_path = "/datasets/DongbeiDaxue/chengkun_remove_only_tumor"

    for i in range(80):
        case_id = "case_{}".format(str(i).zfill(5))
        case_path = os.path.join(data_path, case_id)
        segmentation_path = os.path.join(case_path, "segmentation")
        image_path = os.path.join(case_path, "imaging")

        destination_patient_path = os.path.join(save_path, "case_{}".format(str(i).zfill(5)))
        destination_patient_path = Path(destination_patient_path)
        destination_image_path = destination_patient_path / "imaging"
        destination_mask_path = destination_patient_path / "segmentation"
        print(destination_image_path, destination_mask_path)
        if not destination_image_path.exists():
            destination_image_path.mkdir(parents=True)
        if not destination_mask_path.exists():
            destination_mask_path.mkdir(parents=True)

        j = 0
        for item in sorted(os.listdir(segmentation_path)):
            segmentation_file = os.path.join(segmentation_path, item)
            segmentation = np.load(segmentation_file)
            image_file = os.path.join(image_path, item)
            image = np.load(image_file)
            if np.max(segmentation) <= 1:
                continue

            np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(j).zfill(3))), image)
            np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(j).zfill(3))), segmentation)
            j += 1

def save_only_liver():
    data_path = "/datasets/DongbeiDaxue/chengkun_remove"
    save_path = "/datasets/DongbeiDaxue/chengkun_remove_only_liver"

    for i in range(80):
        case_id = "case_{}".format(str(i).zfill(5))
        case_path = os.path.join(data_path, case_id)
        segmentation_path = os.path.join(case_path, "segmentation")
        image_path = os.path.join(case_path, "imaging")

        destination_patient_path = os.path.join(save_path, "case_{}".format(str(i).zfill(5)))
        destination_patient_path = Path(destination_patient_path)
        destination_image_path = destination_patient_path / "imaging"
        destination_mask_path = destination_patient_path / "segmentation"
        print(destination_image_path, destination_mask_path)
        if not destination_image_path.exists():
            destination_image_path.mkdir(parents=True)
        if not destination_mask_path.exists():
            destination_mask_path.mkdir(parents=True)

        j = 0
        for item in sorted(os.listdir(segmentation_path)):
            segmentation_file = os.path.join(segmentation_path, item)
            segmentation = np.load(segmentation_file)
            image_file = os.path.join(image_path, item)
            image = np.load(image_file)
            if np.max(segmentation) <= 0:
                continue

            np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(j).zfill(3))), image)
            np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(j).zfill(3))), segmentation)
            j += 1

#对肿瘤数据进行分析
def analysis_only_tumor():
    data_path = "/datasets/DongbeiDaxue/chengkun_remove_only_tumor"
    max_tumor = 0
    tumor_size = 0
    background_size = 0
    liver_size = 0

    for i in range(80):
        case_id = "case_{}".format(str(i).zfill(5))
        case_path = os.path.join(data_path, case_id)
        segmentation_path = os.path.join(case_path, "segmentation")
        image_path = os.path.join(case_path, "imaging")
        np_list = []

        for item in sorted(os.listdir(segmentation_path)):
            item_file = os.path.join(segmentation_path, item)
            data = np.load(item_file)
            np_list.append(data)

        data = np.array(np_list)

        tumor_size += (data >= 2).sum()
        background_size += (data == 0).sum()
        liver_size += (data == 1).sum()

        if data is None or len(data) == 0:
            continue

        desp = []
        for j in range(2, np.max(data)+1):
            desp.append("{}:{}".format(j, (data == j).sum()))
        print("{}\t{}\t".format(i, ",".join(desp)), (data >= 2).sum())

    print(max_tumor, background_size, liver_size, tumor_size)

#针对数据进行比对
def analysis_tumor():
    tumor = []
    with open("only_tumor_analysis", "r") as file:
        for line in file:
            data = line.strip().split("\t")[1]
            for item in data.split(","):
                tumor.append(int(item.split(":")[1]))

    size = len(tumor)
    print(size)
    tumor = np.array(tumor)
    avg_tumor = np.sum(tumor)/len(tumor)
    print(avg_tumor, np.max(tumor))

    big = 0
    middle = 0
    small = 0
    for item in tumor:
        if item >= 100000:
            big += 1
        elif item >= 10000:
            middle += 1
        else:
            small += 1

    print(big, big/size)
    print(middle, middle/size)
    print(small, small/size)

if __name__ == '__main__':
    save_only_liver()